from flask import Flask, request, jsonify
from flask_cors import CORS
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_ollama import ChatOllama
from langgraph.prebuilt import create_react_agent
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
import os
import socket

# 初始化配置
os.environ["TAVILY_API_KEY"] = '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'
os.environ["DASHSCOPE_API_KEY"] = 'sk-5bc0688c5761427cadb9df012e589136'

# 创建大模型和数据库连接（全局对象）
llm = ChatTongyi(model="qwen-max")
llm_deepseek = ChatOllama(base_url="https://margin-hitachi-combinations-literally.trycloudflare.com", model="deepseek-r1:14b")  # 使用正确的 URL

db = SQLDatabase.from_uri('mysql+pymysql://root:123456@127.0.0.1:3306/mybatis_plus')
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

# 系统提示词
system_prompt = """
您是一个被设计用来与洁客后台系统与数据库交互的智能助手，你的名字叫洁客超级智能体，设计你的人是研发组
的工程师令狐。
您的任务是：
1. **分析用户问题**，并创建一个 **正确的 SQL 查询**。
2. **执行 SQL 查询** 并返回相应的 **查询结果**。
3. **解读查询结果**，用 **自然语言** 向用户反馈答案。
4. **限制查询数据量**，除非用户明确指定，否则默认最多返回 10 条数据。
5. 允许你对数据库执行操作（INSERT、UPDATE、DELETE 等），操作成功以后要提醒我你操作成功了。
6. **先查看数据库中的表结构**，然后再执行查询，以确保 SQL 语句的正确性。
7.允许你
"""
# **不得执行任何修改数据库的操作（INSERT、UPDATE、DELETE 等）**。
# 创建 agent（全局对象）
prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content=system_prompt),
    MessagesPlaceholder(variable_name="messages")
])
agent_executor = create_react_agent(llm, tools, prompt=prompt)

# 创建 Flask 应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 获取本地IP地址
def get_local_ips():
    ips = set()
    try:
        # 获取通过外部连接的IP
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(('8.8.8.8', 1))
            ips.add(s.getsockname()[0])
    except:
        pass
    try:
        # 获取主机名对应的IP
        hostname = socket.gethostname()
        ips.add(socket.gethostbyname(hostname))
    except:
        pass
    # 获取所有接口的IP
    try:
        for interface in socket.getaddrinfo(hostname, None):
            ip = interface[4][0]
            family = interface[0]
            if family == socket.AF_INET:
                ips.add(ip)
    except:
        pass
    ips.discard('127.0.0.1')
    return ['127.0.0.1'] + list(ips)

local_ips = get_local_ips()

# 构造启动日志信息
startup_info = [
    f"* Serving AI app '{app.name}'",
    f"* Debug mode: {'on' if app.debug else 'off'}",
    *[f"* Running on http://{ip}:5000" for ip in local_ips],
    ">>>===JiekeeAI智能体构建成功===<<<",
    "* Restarting with stat",
    "* Debugger is active!",
    "* Designer BY Linghu",
    "* [====内部团队使用======]",
    "* [====版本V1.0======]",
    "* Debugger PIN: [动态生成，启动时可见]"
]


def is_database_query(question):
    # 这里简单地通过关键词判断是否为数据库查询，实际应用中可以更加复杂
    keywords = ['查询', '数据', '表', '记录']
    return any(keyword in question for keyword in keywords)

@app.route('/query', methods=['POST'])
def handle_query():
    data = request.json
    if not data or 'content' not in data:
        return jsonify({"error": "Missing 'content' parameter"}), 400

    question = data['content']

    if is_database_query(question):
        # 使用 ChatTongyi 和 agent_executor 处理数据库查询
        try:
            resp = agent_executor.invoke({
                "messages": [HumanMessage(content=question)]
            })
            result = resp['messages']
            return jsonify({
                "answer": result[-1].content,
                "full_messages": [msg.dict() for msg in result],
                "startup_logs": startup_info
            })
        except Exception as e:
            return jsonify({"error": str(e)}), 500
    else:
        # 处理非数据库相关问题
        try:
            response = llm_deepseek.invoke(question)
            # 直接提取Langchain AIMessage的content
            return jsonify({
                "answer": response.content,
                "startup_logs": startup_info
            })
        except Exception as e:
            return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)